// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <arrow/c/abi.h>
#include <arrow/c/bridge.h>
#include <arrow/testing/gtest_util.h>
#include <arrow/util/config.h>
#include <gtest/gtest.h>

#include <pollux/testing/gtest_utils.h>
#include <pollux/vector/arrow/bridge.h>

namespace kumo::pollux::test {
    namespace {
        static void mockRelease(ArrowSchema *) {
        }

        class ArrowBridgeSchemaExportTest : public testing::Test {
        protected:
            static void SetUpTestCase() {
                memory::MemoryManager::testingSetInstance({});
            }

            void testScalarType(
                const TypePtr &type,
                const char *arrowFormat,
                const ArrowOptions &options = ArrowOptions{}) {
                ArrowSchema arrowSchema;
                export_to_arrow(type, arrowSchema, options);

                verifyScalarType(arrowSchema, arrowFormat);

                arrowSchema.release(&arrowSchema);
                EXPECT_EQ(nullptr, arrowSchema.release);
                EXPECT_EQ(nullptr, arrowSchema.private_data);
            }

            void verifyScalarType(
                const ArrowSchema &arrowSchema,
                const char *arrowFormat,
                const char *name = nullptr) {
                EXPECT_STREQ(arrowFormat, arrowSchema.format);
                if (name == nullptr) {
                    EXPECT_EQ(nullptr, arrowSchema.name);
                } else {
                    EXPECT_STREQ(name, arrowSchema.name);
                }
                EXPECT_EQ(nullptr, arrowSchema.metadata);
                EXPECT_EQ(arrowSchema.flags | ARROW_FLAG_NULLABLE, ARROW_FLAG_NULLABLE);

                EXPECT_EQ(0, arrowSchema.n_children);
                EXPECT_EQ(nullptr, arrowSchema.children);
                EXPECT_EQ(nullptr, arrowSchema.dictionary);
                EXPECT_NE(nullptr, arrowSchema.release);
            }

            // Doesn't check the actual format string of the scalar leaf types (this is
            // tested by the function above), but tests that the types are nested in the
            // correct way.
            void testNestedType(const TypePtr &type) {
                ArrowSchema arrowSchema;
                export_to_arrow(type, arrowSchema);

                verifyNestedType(type, &arrowSchema);

                arrowSchema.release(&arrowSchema);
                EXPECT_EQ(nullptr, arrowSchema.release);
                EXPECT_EQ(nullptr, arrowSchema.private_data);
            }

            void verifyNestedType(const TypePtr &type, ArrowSchema *schema) {
                if (type->kind() == TypeKind::ARRAY) {
                    EXPECT_STREQ("+l", schema->format);
                } else if (type->kind() == TypeKind::MAP) {
                    EXPECT_STREQ("+m", schema->format);
                    ASSERT_EQ(schema->n_children, 1);
                    schema = schema->children[0];
                    // Map data should be a non-nullable struct type
                    ASSERT_EQ(schema->flags & ARROW_FLAG_NULLABLE, 0);
                    ASSERT_EQ(schema->n_children, 2);
                    // Map data key type should be a non-nullable
                    ASSERT_EQ(schema->children[0]->flags & ARROW_FLAG_NULLABLE, 0);
                } else if (type->kind() == TypeKind::ROW) {
                    EXPECT_STREQ("+s", schema->format);
                }
                // Scalar type.
                else {
                    EXPECT_EQ(nullptr, schema->children);
                }
                ASSERT_EQ(type->size(), schema->n_children);

                // Recurse down the children.
                for (size_t i = 0; i < type->size(); ++i) {
                    verifyNestedType(type->childAt(i), schema->children[i]);

                    // If this is a rowType, assert that the children returned with the
                    // correct name set.
                    if (auto rowType = std::dynamic_pointer_cast<const RowType>(type)) {
                        EXPECT_EQ(rowType->nameOf(i), std::string(schema->children[i]->name));
                    }
                }
            }

            void testConstant(
                const TypePtr &type,
                const char *arrowFormat,
                const ArrowOptions &options = ArrowOptions{}) {
                ArrowSchema arrowSchema;
                const bool is_scalar = (type->size() == 0);
                const bool constantSize = 100;

                // If scalar, create the constant vector directly; if complex type, create a
                // complex vector first, then wrap it in a dictionary.
                auto constant_vector = is_scalar
                                           ? BaseVector::create_constant(
                                               type, variant(type->kind()), constantSize, pool_.get())
                                           : BaseVector::wrap_in_constant(
                                               constantSize,
                                               3, // index to use for the constant
                                               BaseVector::create(type, 100, pool_.get()));

                pollux::export_to_arrow(constant_vector, arrowSchema, options);

                EXPECT_STREQ("+r", arrowSchema.format);
                EXPECT_EQ(nullptr, arrowSchema.name);

                EXPECT_EQ(2, arrowSchema.n_children);
                EXPECT_NE(nullptr, arrowSchema.children);
                EXPECT_EQ(nullptr, arrowSchema.dictionary);

                // Validate run_ends.
                EXPECT_NE(nullptr, arrowSchema.children[0]);
                const auto &runEnds = *arrowSchema.children[0];

                EXPECT_STREQ("i", runEnds.format);
                EXPECT_STREQ("run_ends", runEnds.name);
                EXPECT_EQ(0, runEnds.n_children);
                EXPECT_EQ(nullptr, runEnds.children);
                EXPECT_EQ(nullptr, runEnds.dictionary);

                // Validate values.
                EXPECT_NE(nullptr, arrowSchema.children[1]);

                if (is_scalar) {
                    verifyScalarType(*arrowSchema.children[1], arrowFormat, "values");
                } else {
                    EXPECT_STREQ(arrowFormat, arrowSchema.children[1]->format);
                    verifyNestedType(type, arrowSchema.children[1]);
                }

                arrowSchema.release(&arrowSchema);
                EXPECT_EQ(nullptr, arrowSchema.release);
                EXPECT_EQ(nullptr, arrowSchema.private_data);
            }

            void export_to_arrow(
                const TypePtr &type,
                ArrowSchema &out,
                const ArrowOptions &options = ArrowOptions{}) {
                pollux::export_to_arrow(
                    BaseVector::create(type, 0, pool_.get()), out, options);
            }

            ArrowSchema makeArrowSchema(const char *format) {
                return ArrowSchema{
                    .format = format,
                    .name = nullptr,
                    .metadata = nullptr,
                    .flags = 0,
                    .n_children = 0,
                    .children = nullptr,
                    .dictionary = nullptr,
                    .release = mockRelease,
                    .private_data = nullptr,
                };
            }

            std::shared_ptr<memory::MemoryPool> pool_{
                memory::memoryManager()->addLeafPool()
            };
        };

        TEST_F(ArrowBridgeSchemaExportTest, scalar) {
            testScalarType(TINYINT(), "c");
            testScalarType(SMALLINT(), "s");
            testScalarType(INTEGER(), "i");
            testScalarType(BIGINT(), "l");

            testScalarType(BOOLEAN(), "b");

            testScalarType(REAL(), "f");
            testScalarType(DOUBLE(), "g");

            testScalarType(VARCHAR(), "u");
            testScalarType(VARBINARY(), "z");

            // Test default timezone
            testScalarType(
                TIMESTAMP(), "tss:", {.timestampUnit = TimestampUnit::kSecond});
            testScalarType(TIMESTAMP(), "tsm:", {.timestampUnit = TimestampUnit::kMilli});
            testScalarType(TIMESTAMP(), "tsu:", {.timestampUnit = TimestampUnit::kMicro});
            testScalarType(TIMESTAMP(), "tsn:", {.timestampUnit = TimestampUnit::kNano});

            testScalarType(VARCHAR(), "vu", {.exportToStringView = true});
            testScalarType(VARBINARY(), "vz", {.exportToStringView = true});

            // Test specific timezone
            testScalarType(
                TIMESTAMP(),
                "tss:+01:0",
                {.timestampUnit = TimestampUnit::kSecond, .timestampTimeZone = "+01:0"});
            testScalarType(
                TIMESTAMP(),
                "tsm:+01:0",
                {.timestampUnit = TimestampUnit::kMilli, .timestampTimeZone = "+01:0"});
            testScalarType(
                TIMESTAMP(),
                "tsu:+01:0",
                {.timestampUnit = TimestampUnit::kMicro, .timestampTimeZone = "+01:0"});
            testScalarType(
                TIMESTAMP(),
                "tsn:+01:0",
                {.timestampUnit = TimestampUnit::kNano, .timestampTimeZone = "+01:0"});

            testScalarType(DATE(), "tdD");
            testScalarType(INTERVAL_YEAR_MONTH(), "tiM");

            testScalarType(DECIMAL(10, 4), "d:10,4");
            testScalarType(DECIMAL(20, 15), "d:20,15");

            testScalarType(UNKNOWN(), "n");
        }

        TEST_F(ArrowBridgeSchemaExportTest, nested) {
            // Array.
            testNestedType(ARRAY(INTEGER()));
            testNestedType(ARRAY(VARCHAR()));
            testNestedType(ARRAY(ARRAY(TINYINT())));
            testNestedType(ARRAY(ARRAY(ARRAY(ARRAY(BOOLEAN())))));

            // Map.
            testNestedType(MAP(INTEGER(), DOUBLE()));
            testNestedType(MAP(VARBINARY(), BOOLEAN()));
            testNestedType(MAP(VARBINARY(), MAP(SMALLINT(), REAL())));
            testNestedType(MAP(VARBINARY(), MAP(SMALLINT(), MAP(INTEGER(), BIGINT()))));

            // Row.
            testNestedType(ROW({}));
            testNestedType(ROW({INTEGER()}));
            testNestedType(ROW({INTEGER(), DOUBLE()}));
            testNestedType(
                ROW({INTEGER(), DOUBLE(), ROW({BIGINT(), REAL(), BOOLEAN()})}));

            // Row with names.
            testNestedType(ROW({"my_col"}, {INTEGER()}));
            testNestedType(ROW({"my_col", "my_other_col"}, {INTEGER(), VARCHAR()}));

            // Mix and match.
            testNestedType(
                ROW({"c1", "c2", "c3"},
                    {
                        ARRAY(INTEGER()),
                        MAP(ROW({VARBINARY(), SMALLINT()}), BOOLEAN()),
                        ARRAY(MAP(INTEGER(), VARCHAR())),
                    }));
        }

        TEST_F(ArrowBridgeSchemaExportTest, constant) {
            testConstant(TINYINT(), "c");
            testConstant(INTEGER(), "i");
            testConstant(BOOLEAN(), "b");
            testConstant(DOUBLE(), "g");
            testConstant(VARCHAR(), "u");
            testConstant(VARCHAR(), "vu", {.exportToStringView = true});
            testConstant(DATE(), "tdD");
            testConstant(INTERVAL_YEAR_MONTH(), "tiM");
            testConstant(UNKNOWN(), "n");

            testConstant(ARRAY(INTEGER()), "+l");
            testConstant(ARRAY(UNKNOWN()), "+l");
            testConstant(MAP(BOOLEAN(), REAL()), "+m");
            testConstant(MAP(UNKNOWN(), REAL()), "+m");
            testConstant(ROW({TIMESTAMP(), DOUBLE()}), "+s");
            testConstant(ROW({UNKNOWN(), UNKNOWN()}), "+s");
            POLLUX_ASSERT_THROW(
                testConstant(ARRAY(INTEGER()), "+l", {false, true}),
                "Flattening is only supported for scalar types.");
            POLLUX_ASSERT_THROW(
                testConstant(MAP(BOOLEAN(), REAL()), "+m", {false, true}),
                "Flattening is only supported for scalar types.");
            POLLUX_ASSERT_THROW(
                testConstant(ROW({BOOLEAN(), REAL()}), "+s", {false, true}),
                "Flattening is only supported for scalar types.");
        }

        class ArrowBridgeSchemaImportTest : public ArrowBridgeSchemaExportTest {
        protected:
            TypePtr testSchemaImport(const char *format) {
                auto arrowSchema = makeArrowSchema(format);
                auto type = import_from_arrow(arrowSchema);
                arrowSchema.release(&arrowSchema);
                return type;
            }

            TypePtr testSchemaDictionaryImport(const char *indexFmt, ArrowSchema schema) {
                auto dictionarySchema = makeArrowSchema(indexFmt);
                dictionarySchema.dictionary = &schema;

                auto type = import_from_arrow(dictionarySchema);
                dictionarySchema.release(&dictionarySchema);
                return type;
            }

            TypePtr testSchemaReeImport(const char *valuesFmt) {
                auto reeSchema = makeArrowSchema("+r");
                auto runsSchema = makeArrowSchema("i");
                auto valuesSchema = makeArrowSchema(valuesFmt);

                std::vector<ArrowSchema *> schemas{&runsSchema, &valuesSchema};
                reeSchema.n_children = 2;
                reeSchema.children = schemas.data();

                auto type = import_from_arrow(reeSchema);
                reeSchema.release(&reeSchema);
                return type;
            }

            ArrowSchema makeComplexArrowSchema(
                std::vector<ArrowSchema> &schemas,
                std::vector<ArrowSchema *> &schemaPtrs,
                std::vector<ArrowSchema> &mapSchemas,
                std::vector<ArrowSchema *> &mapSchemaPtrs,
                const char *mainFormat,
                const std::vector<const char *> &childrenFormat,
                const std::vector<const char *> &colNames = {}) {
                schemas.clear();
                schemaPtrs.clear();
                mapSchemas.clear();
                mapSchemaPtrs.clear();
                schemas.resize(childrenFormat.size());
                schemaPtrs.resize(childrenFormat.size());

                for (size_t i = 0; i < childrenFormat.size(); ++i) {
                    schemas[i] = makeArrowSchema(childrenFormat[i]);
                    if (colNames.size() > i) {
                        schemas[i].name = colNames[i];
                    }
                    schemaPtrs[i] = &schemas[i];
                }

                auto mainSchema = makeArrowSchema(mainFormat);
                if (strcmp(mainFormat, "+m") == 0) {
                    // Arrow wraps key and value in a struct.
                    mapSchemas.resize(1);
                    mapSchemaPtrs.resize(1);
                    mapSchemas[0] = makeArrowSchema("+s");
                    auto *child = &mapSchemas[0];
                    mapSchemaPtrs[0] = &mapSchemas[0];
                    child->n_children = schemaPtrs.size();
                    child->children = schemaPtrs.data();
                    mainSchema.n_children = 1;
                    mainSchema.children = mapSchemaPtrs.data();
                } else {
                    mainSchema.n_children = (int64_t) schemaPtrs.size();
                    mainSchema.children = schemaPtrs.data();
                }

                return mainSchema;
            }

            TypePtr testSchemaImportComplex(
                const char *mainFormat,
                const std::vector<const char *> &childrenFormat,
                const std::vector<const char *> &colNames = {}) {
                std::vector<ArrowSchema> schemas;
                std::vector<ArrowSchema> mapSchemas;
                std::vector<ArrowSchema *> schemaPtrs;
                std::vector<ArrowSchema *> mapSchemaPtrs;
                auto type = import_from_arrow(makeComplexArrowSchema(
                    schemas,
                    schemaPtrs,
                    mapSchemas,
                    mapSchemaPtrs,
                    mainFormat,
                    childrenFormat,
                    colNames));
                return type;
            }
        };

        TEST_F(ArrowBridgeSchemaImportTest, scalar) {
            EXPECT_EQ(*BOOLEAN(), *testSchemaImport("b"));
            EXPECT_EQ(*TINYINT(), *testSchemaImport("c"));
            EXPECT_EQ(*SMALLINT(), *testSchemaImport("s"));
            EXPECT_EQ(*INTEGER(), *testSchemaImport("i"));
            EXPECT_EQ(*BIGINT(), *testSchemaImport("l"));
            EXPECT_EQ(*REAL(), *testSchemaImport("f"));
            EXPECT_EQ(*DOUBLE(), *testSchemaImport("g"));

            EXPECT_EQ(*VARCHAR(), *testSchemaImport("u"));
            EXPECT_EQ(*VARCHAR(), *testSchemaImport("U"));
            EXPECT_EQ(*VARBINARY(), *testSchemaImport("z"));
            EXPECT_EQ(*VARBINARY(), *testSchemaImport("Z"));

            EXPECT_EQ(*VARCHAR(), *testSchemaImport("vu"));
            EXPECT_EQ(*VARBINARY(), *testSchemaImport("vz"));

            // Temporal.
            EXPECT_EQ(*TIMESTAMP(), *testSchemaImport("tsn:"));
            EXPECT_EQ(*DATE(), *testSchemaImport("tdD"));
            EXPECT_EQ(*INTERVAL_YEAR_MONTH(), *testSchemaImport("tiM"));

            EXPECT_EQ(*DECIMAL(10, 4), *testSchemaImport("d:10,4"));
            EXPECT_EQ(*DECIMAL(20, 15), *testSchemaImport("d:20,15"));
            POLLUX_ASSERT_THROW(
                *testSchemaImport("d2,15"),
                "Unable to convert 'd2,15' ArrowSchema decimal format to Pollux decimal");
            EXPECT_EQ(*DECIMAL(10, 4), *testSchemaImport("d:10,4,128"));
            POLLUX_ASSERT_THROW(
                *testSchemaImport("d:10,4,256"),
                "Conversion failed for 'd:10,4,256'. Pollux decimal does not support custom bitwidth.");
            POLLUX_ASSERT_THROW(
                *testSchemaImport("d:10,4,"),
                "Unable to convert 'd:10,4,' ArrowSchema decimal format to Pollux decimal");
            POLLUX_ASSERT_THROW(
                *testSchemaImport("d:10"),
                "Unable to convert 'd:10' ArrowSchema decimal format to Pollux decimal");
            POLLUX_ASSERT_THROW(
                *testSchemaImport("d:"),
                "Unable to convert 'd:' ArrowSchema decimal format to Pollux decimal");
            POLLUX_ASSERT_THROW(
                *testSchemaImport("d:10,"),
                "Unable to convert 'd:10,' ArrowSchema decimal format to Pollux decimal");
        }

        TEST_F(ArrowBridgeSchemaImportTest, complexTypes) {
            // Array.
            EXPECT_EQ(*ARRAY(BIGINT()), *testSchemaImportComplex("+l", {"l"}));
            EXPECT_EQ(*ARRAY(TIMESTAMP()), *testSchemaImportComplex("+l", {"tsn:"}));
            EXPECT_EQ(*ARRAY(DATE()), *testSchemaImportComplex("+l", {"tdD"}));
            EXPECT_EQ(
                *ARRAY(INTERVAL_YEAR_MONTH()), *testSchemaImportComplex("+l", {"tiM"}));
            EXPECT_EQ(*ARRAY(VARCHAR()), *testSchemaImportComplex("+l", {"U"}));

            EXPECT_EQ(*ARRAY(DECIMAL(10, 4)), *testSchemaImportComplex("+l", {"d:10,4"}));
            EXPECT_EQ(
                *ARRAY(DECIMAL(20, 15)), *testSchemaImportComplex("+l", {"d:20,15"}));

            // Map.
            EXPECT_EQ(
                *MAP(VARCHAR(), BOOLEAN()), *testSchemaImportComplex("+m", {"U", "b"}));
            EXPECT_EQ(
                *MAP(SMALLINT(), REAL()), *testSchemaImportComplex("+m", {"s", "f"}));

            // Row/struct.
            EXPECT_EQ(
                *ROW({SMALLINT(), REAL()}), *testSchemaImportComplex("+s", {"s", "f"}));
            EXPECT_EQ(
                *ROW({SMALLINT(), REAL(), VARCHAR(), BOOLEAN()}),
                *testSchemaImportComplex("+s", {"s", "f", "u", "b"}));

            // Named
            EXPECT_EQ(
                *ROW({"col1", "col2"}, {SMALLINT(), REAL()}),
                *testSchemaImportComplex("+s", {"s", "f"}, {"col1", "col2"}));
        }

        TEST_F(ArrowBridgeSchemaImportTest, unsupported) {
            EXPECT_THROW(testSchemaImport("C"), PolluxUserError);
            EXPECT_THROW(testSchemaImport("S"), PolluxUserError);
            EXPECT_THROW(testSchemaImport("I"), PolluxUserError);
            EXPECT_THROW(testSchemaImport("L"), PolluxUserError);
            EXPECT_THROW(testSchemaImport("e"), PolluxUserError);

            EXPECT_THROW(testSchemaImport("w:42"), PolluxUserError);

            EXPECT_THROW(testSchemaImport("tdm"), PolluxUserError);
            EXPECT_THROW(testSchemaImport("tts"), PolluxUserError);
            EXPECT_THROW(testSchemaImport("ttm"), PolluxUserError);
            EXPECT_THROW(testSchemaImport("tDs"), PolluxUserError);

            EXPECT_THROW(testSchemaImport("+"), PolluxUserError);
            EXPECT_THROW(testSchemaImport("+L"), PolluxUserError);
            EXPECT_THROW(testSchemaImport("+b"), PolluxUserError);
            EXPECT_THROW(testSchemaImport("+z"), PolluxUserError);
            EXPECT_THROW(testSchemaImport("+u"), PolluxUserError);
            EXPECT_THROW(testSchemaImport("+w"), PolluxUserError);
        }

        class ArrowBridgeSchemaTest : public testing::Test {
        protected:
            static void SetUpTestCase() {
                memory::MemoryManager::testingSetInstance({});
            }

            void roundtripTest(
                const TypePtr &inputType,
                const ArrowOptions &options = ArrowOptions{}) {
                ArrowSchema arrowSchema;
                export_to_arrow(inputType, arrowSchema, options);
                auto outputType = import_from_arrow(arrowSchema);
                arrowSchema.release(&arrowSchema);
                EXPECT_EQ(*inputType, *outputType);
            }

            void export_to_arrow(
                const TypePtr &type,
                ArrowSchema &out,
                const ArrowOptions &options = ArrowOptions{}) {
                pollux::export_to_arrow(
                    BaseVector::create(type, 0, pool_.get()), out, options);
            }

            std::shared_ptr<memory::MemoryPool> pool_{
                memory::memoryManager()->addLeafPool()
            };
        };

        TEST_F(ArrowBridgeSchemaTest, roundtrip) {
            roundtripTest(BOOLEAN());
            roundtripTest(VARCHAR());
            roundtripTest(VARCHAR(), {.exportToStringView = true});
            roundtripTest(REAL());
            roundtripTest(ARRAY(DOUBLE()));
            roundtripTest(ARRAY(ARRAY(ARRAY(ARRAY(VARBINARY())))));
            roundtripTest(MAP(VARCHAR(), REAL()));
            roundtripTest(MAP(VARCHAR(), ARRAY(BOOLEAN())));
            roundtripTest(MAP(VARCHAR(), ARRAY(MAP(ARRAY(BIGINT()), BOOLEAN()))));
            roundtripTest(ROW({VARBINARY(), TINYINT(), SMALLINT()}));
            roundtripTest(ROW({VARBINARY(), ROW({DOUBLE(), VARBINARY()}), SMALLINT()}));
            roundtripTest(ROW({
                ARRAY(VARBINARY()),
                MAP(REAL(), ARRAY(DOUBLE())),
                ROW({"a", "b"}, {DOUBLE(), VARBINARY()}),
                INTEGER(),
            }));
        }

        TEST_F(ArrowBridgeSchemaTest, validateInArrow) {
            const std::pair<TypePtr, std::shared_ptr<arrow::DataType> > kTypes[] = {
                {BOOLEAN(), arrow::boolean()},
                {VARCHAR(), arrow::utf8()},
                {VARCHAR(), arrow::utf8_view()},
#if ARROW_VERSION_MAJOR >= 18
      {DECIMAL(10, 4), arrow::decimal128(10, 4)},
      {DECIMAL(20, 15), arrow::decimal128(20, 15)},
#else
                {DECIMAL(10, 4), arrow::decimal(10, 4)},
                {DECIMAL(20, 15), arrow::decimal(20, 15)},
#endif
                {ARRAY(DOUBLE()), arrow::list(arrow::float64())},
                {ARRAY(ARRAY(DOUBLE())), arrow::list(arrow::list(arrow::float64()))},
                {MAP(VARCHAR(), REAL()), arrow::map(arrow::utf8(), arrow::float32())},
                {
                    ROW({"c0", "c1"}, {VARCHAR(), BIGINT()}),
                    arrow::struct_(
                        {
                            arrow::field("c0", arrow::utf8()),
                            arrow::field("c1", arrow::int64())
                        })
                },
            };
            for (auto &[tv, ta]: kTypes) {
                VKLOG(1) << "Validating conversion between " << tv->toString() << " and "
            << ta->ToString();
                ArrowSchema schema;
                ta == arrow::utf8_view()
                    ? export_to_arrow(tv, schema, {.exportToStringView = true})
                    : export_to_arrow(tv, schema);
                ASSERT_OK_AND_ASSIGN(auto actual, arrow::ImportType(&schema));
                ASSERT_FALSE(schema.release);
                EXPECT_EQ(*actual, *ta);
            }
        }

        TEST_F(ArrowBridgeSchemaImportTest, dictionaryTypeTest) {
            // Primitive types
            EXPECT_EQ(DOUBLE(), testSchemaDictionaryImport("i", makeArrowSchema("g")));
            EXPECT_EQ(BOOLEAN(), testSchemaDictionaryImport("i", makeArrowSchema("b")));
            EXPECT_EQ(TINYINT(), testSchemaDictionaryImport("i", makeArrowSchema("c")));
            EXPECT_EQ(INTEGER(), testSchemaDictionaryImport("i", makeArrowSchema("i")));
            EXPECT_EQ(SMALLINT(), testSchemaDictionaryImport("i", makeArrowSchema("s")));
            EXPECT_EQ(BIGINT(), testSchemaDictionaryImport("i", makeArrowSchema("l")));
            EXPECT_EQ(REAL(), testSchemaDictionaryImport("i", makeArrowSchema("f")));
            EXPECT_EQ(VARCHAR(), testSchemaDictionaryImport("i", makeArrowSchema("u")));

            std::vector<ArrowSchema> schemas;
            std::vector<ArrowSchema> mapSchemas;
            std::vector<ArrowSchema *> mapSchemaPtrs;
            std::vector<ArrowSchema *> schemaPtrs;

            // Arrays
            EXPECT_EQ(
                *ARRAY(BIGINT()),
                *testSchemaDictionaryImport(
                    "i",
                    makeComplexArrowSchema(
                        schemas, schemaPtrs, mapSchemas, mapSchemaPtrs, "+l", {"l"})));
            EXPECT_EQ(
                *ARRAY(TIMESTAMP()),
                *testSchemaDictionaryImport(
                    "i",
                    makeComplexArrowSchema(
                        schemas, schemaPtrs, mapSchemas, mapSchemaPtrs, "+l", {"tsn:"})));
            EXPECT_EQ(
                *ARRAY(DATE()),
                *testSchemaDictionaryImport(
                    "i",
                    makeComplexArrowSchema(
                        schemas, schemaPtrs, mapSchemas, mapSchemaPtrs, "+l", {"tdD"})));
            EXPECT_EQ(
                *ARRAY(INTERVAL_YEAR_MONTH()),
                *testSchemaDictionaryImport(
                    "i",
                    makeComplexArrowSchema(
                        schemas, schemaPtrs, mapSchemas, mapSchemaPtrs, "+l", {"tiM"})));
            EXPECT_EQ(
                *ARRAY(VARCHAR()),
                *testSchemaDictionaryImport(
                    "i",
                    makeComplexArrowSchema(
                        schemas, schemaPtrs, mapSchemas, mapSchemaPtrs, "+l", {"U"})));

            // Maps
            EXPECT_EQ(
                *MAP(VARCHAR(), BOOLEAN()),
                *testSchemaDictionaryImport(
                    "i",
                    makeComplexArrowSchema(
                        schemas,
                        schemaPtrs,
                        mapSchemas,
                        mapSchemaPtrs,
                        "+m",
                        {"U", "b"})));
            EXPECT_EQ(
                *MAP(SMALLINT(), REAL()),
                *testSchemaDictionaryImport(
                    "i",
                    makeComplexArrowSchema(
                        schemas,
                        schemaPtrs,
                        mapSchemas,
                        mapSchemaPtrs,
                        "+m",
                        {"s", "f"})));

            // Rows
            EXPECT_EQ(
                *ROW({SMALLINT(), REAL()}),
                *testSchemaDictionaryImport(
                    "i",
                    makeComplexArrowSchema(
                        schemas,
                        schemaPtrs,
                        mapSchemas,
                        mapSchemaPtrs,
                        "+s",
                        {"s", "f"})));

            // Named Row
            EXPECT_EQ(
                *ROW({"col1", "col2"}, {SMALLINT(), REAL()}),
                *testSchemaDictionaryImport(
                    "i",
                    makeComplexArrowSchema(
                        schemas,
                        schemaPtrs,
                        mapSchemas,
                        mapSchemaPtrs,
                        "+s",
                        {"s", "f"},
                        {"col1", "col2"})));
        }

        TEST_F(ArrowBridgeSchemaImportTest, reeTypeTest) {
            // Ensure REE just returns the type of the inner `values` child.
            EXPECT_EQ(DOUBLE(), testSchemaReeImport("g"));
            EXPECT_EQ(INTEGER(), testSchemaReeImport("i"));
            EXPECT_EQ(BIGINT(), testSchemaReeImport("l"));
        }
    } // namespace
} // namespace kumo::pollux::test
